import os

import SimpleITK as sitk
import numpy as np
from scipy import ndimage
from skimage import measure
from tqdm import tqdm
from Utils import save_nii_file

smooth = 1.


def save_itk(image, origin, spacing, filename):
    if type(origin) != tuple:
        if type(origin) == list:
            origin = tuple(reversed(origin))
        else:
            origin = tuple(reversed(origin.tolist()))
    if type(spacing) != tuple:
        if type(spacing) == list:
            spacing = tuple(reversed(spacing))
        else:
            spacing = tuple(reversed(spacing.tolist()))
    itkimage = sitk.GetImageFromArray(image, isVector=False)
    itkimage.SetSpacing(spacing)
    itkimage.SetOrigin(origin)
    sitk.WriteImage(itkimage, filename, True)


def load_itk_image(filename):
    itkimage = sitk.ReadImage(filename)
    numpyImage = sitk.GetArrayFromImage(itkimage)
    # numpyImage = numpyImage - np.min(numpyImage)
    numpyOrigin = list(reversed(itkimage.GetOrigin()))
    numpySpacing = list(reversed(itkimage.GetSpacing()))
    return numpyImage, numpyOrigin, numpySpacing


def fill_2d_hole(bw):
    for i in range(bw.shape[0]):
        current_slice = bw[i]
        label = measure.label(current_slice, connectivity=2)
        properties = measure.regionprops(label)
        for prop in properties:
            bb = prop.bbox
            current_slice[bb[0]:bb[2], bb[1]:bb[3]] = current_slice[bb[0]:bb[2], bb[1]:bb[3]] | prop.filled_image
        bw[i] = current_slice
    return bw


def fill_hole(curpred, bw, bw_filled):
    """

    :param curpred: curperd  =
    :param bw:
    :param bw_filled:
    :return:
    """
    holes = np.logical_xor(bw, bw_filled)  # 逻辑异或
    save_nii_file(file_path, holes, "holes")
    label = measure.label(holes, connectivity=3)
    save_nii_file(file_path, label, "holes_label")
    properties = measure.regionprops(label)
    struct = ndimage.generate_binary_structure(3, 3)

    for idx in tqdm(range(len(properties))):
        hole = (label == properties[idx].label)
        bb = properties[idx].bbox
        bw0 = bw[bb[0]:bb[3], bb[1]:bb[4], bb[2]:bb[5]]
        # dilation - original
        bw1 = holes[bb[0]:bb[3], bb[1]:bb[4], bb[2]:bb[5]]
        pred0 = curpred[bb[0]:bb[3], bb[1]:bb[4], bb[2]:bb[5]]
        bw2 = ndimage.binary_dilation(bw1, structure=struct, iterations=1)
        # assert(bw1.shape[idx] == bw2.shape[idx] for idx in range(3))
        bw3 = np.logical_xor(bw1, bw2)
        bw4 = np.logical_and(bw3, bw0)
        target = pred0[bw4].flatten()
        if len(target) > 0:
            target = target - np.min(target)  # Make sure the value is positive
            counts = np.bincount(target)
            new_label = np.argmax(counts)
            curpred[hole] = new_label
    return curpred


def bw_dilation(bw, iternum=5):
    bw_org = np.copy(bw)
    struct = ndimage.generate_binary_structure(3, 3)
    bw_dillated = ndimage.binary_dilation(bw_org, structure=struct, iterations=iternum)
    return bw_dillated


def two_lung_only(sample_file_path, bw0, spacing, max_ratio=4.8):
    label = measure.label(bw0, connectivity=3)  # label function is used to connect the connectable region
    properties = measure.regionprops(label)
    properties.sort(key=lambda x: x.area, reverse=True)

    print(properties[0].area)
    # bw1 = (label == properties[0].label)
    # save_nii_file(sample_file_path, bw1, "bw1")
    print(properties[1].area)
    # bw2 = (label == properties[1].label)
    # save_nii_file(sample_file_path, bw2, "bw2")
    if len(properties) > 1 and properties[0].area / properties[1].area < max_ratio:
        bw1 = (label == properties[0].label)
        bw2 = (label == properties[1].label)

        d1 = ndimage.morphology.distance_transform_edt(bw1 == 0, sampling=spacing)
        d2 = ndimage.morphology.distance_transform_edt(bw2 == 0, sampling=spacing)
        bw1 = bw0 & (d1 < d2)
        bw2 = bw0 & (d1 > d2)
        bw = bw1 | bw2

        label2 = measure.label(bw, connectivity=3)
        properties2 = measure.regionprops(label2)
        properties2.sort(key=lambda x: x.area, reverse=True)
        if len(properties2) > 1 and properties2[0].area / properties2[1].area < max_ratio:
            bw1s = (label2 == properties2[0].label)
            bw2s = (label2 == properties2[1].label)
            bw = bw1s | bw2s

            bw_filled = bw.copy()
            bb1 = properties2[0].bbox
            bb2 = properties2[1].bbox
            bw_filled[bb1[0]:bb1[3], bb1[1]:bb1[4], bb1[2]:bb1[5]] = bw_filled[bb1[0]:bb1[3], bb1[1]:bb1[4],
                                                                     bb1[2]:bb1[5]] | properties2[0].filled_image
            bw_filled[bb2[0]:bb2[3], bb2[1]:bb2[4], bb2[2]:bb2[5]] = bw_filled[bb2[0]:bb2[3], bb2[1]:bb2[4],
                                                                     bb2[2]:bb2[5]] | properties2[1].filled_image

        else:
            bw1s = (label2 == properties2[0].label)
            bw = bw1s

            bw_filled = bw.copy()
            bb1 = properties2[0].bbox
            bw_filled[bb1[0]:bb1[3], bb1[1]:bb1[4], bb1[2]:bb1[5]] = bw_filled[bb1[0]:bb1[3], bb1[1]:bb1[4],
                                                                     bb1[2]:bb1[5]] | properties2[0].filled_image

    else:
        current_property = properties[1]
        # Because the maximum area is not the lung, so I should select the second property
        bw1 = (label == current_property.label)
        # save to nii image
        save_nii_file(sample_file_path, bw1, "bw1 or curpred_bw")

        bw = bw1

        bw_filled = bw.copy()
        bb1 = current_property.bbox
        # Actually, bw_filled is the copy of bw1.  bw_filled = bw_filled | properties[0].filled_image
        bw_filled[bb1[0]:bb1[3], bb1[1]:bb1[4], bb1[2]:bb1[5]] = bw_filled[bb1[0]:bb1[3], bb1[1]:bb1[4],
                                                                 bb1[2]:bb1[5]] | current_property.filled_image
        save_nii_file(sample_file_path, bw_filled, "bw_filled or curpred_bw_filled")
    return bw, bw_filled


def extract_largest(curpred, bw_pred):
    curped_new = np.zeros(curpred.shape)
    struct = ndimage.generate_binary_structure(3, 3)
    for i in range(1, 6):
        bwall = (curpred == i)
        label = measure.label(bwall, connectivity=3)
        properties = measure.regionprops(label)
        properties.sort(key=lambda x: x.area, reverse=True)

        for idx in tqdm(range(len(properties))):
            bw = (label == properties[idx].label)
            if idx == 0:
                curped_new[bw] = i
            else:
                bb = properties[idx].bbox
                bw0 = bw_pred[bb[0]:bb[3], bb[1]:bb[4], bb[2]:bb[5]]
                # dilation - original
                bw1 = bw[bb[0]:bb[3], bb[1]:bb[4], bb[2]:bb[5]]
                pred0 = curpred[bb[0]:bb[3], bb[1]:bb[4], bb[2]:bb[5]]
                bw2 = ndimage.binary_dilation(bw1, structure=struct, iterations=1)
                # assert(bw1.shape[idx] == bw2.shape[idx] for idx in range(3))
                bw3 = np.logical_xor(bw1, bw2)
                bw4 = np.logical_and(bw3, bw0)
                target = pred0[bw4].flatten()
                if len(target) > 0:
                    target = target - np.min(target)  # Make sure the value is positive
                    counts = np.bincount(target)
                    new_label = np.argmax(counts)
                    curped_new[bw] = new_label

    return curped_new.astype("uint8")


def postprocess_lung(curpred, spacing):
    # time_start = time.time()
    # optimal_threshold = get_optimal_threshold(curpred)
    optimal_threshold = -359
    curpred_bw = (curpred < optimal_threshold)
    save_nii_file(file_path, curpred_bw, "curpred_bw_just_threshold")
    curpred_bw, curpred_bw_filled = two_lung_only(file_path, curpred_bw, spacing, max_ratio=2)
    curpred_bw = curpred_bw.astype('uint8')
    curpred *= curpred_bw
    save_nii_file(file_path, curpred, "curpred")
    curpred_fill_hole = fill_hole(curpred, curpred_bw, curpred_bw_filled)
    save_nii_file(file_path, curpred_fill_hole.astype("uint8"), "curpred_fill_hole")
    curpred_extract_largest = extract_largest(curpred_fill_hole, curpred_bw_filled)
    save_nii_file(file_path, curpred_extract_largest, "curpred_extract_largest")
    # time_end = time.time()
    # print('elapse {} seconds'.format(time_end-time_start))
    return curpred_extract_largest, curpred_bw


if __name__ == '__main__':
    # RUL 1 RML 2 RLL 3
    # LUL 4 LLL 5
    file_path = r"D:\dataset\vessel\Parse_2022_train_data\PA000005\image\PA000005.nii.gz"
    save_path = os.path.dirname(file_path)
    image, origin, spacing = load_itk_image(file_path)
    print(np.max(image))
    print(np.min(image))
    curpred = postprocess_lung(image, spacing)
    save_itk(curpred.astype('uint8'), origin, spacing, save_path)
    # print("CurID {} Done".format(curID))
